import numpy as np
#from matplotlib.pyplot import plot, savefig, title, legend, ylim, cla, xlabel, ylabel, annotate
from time import clock

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from matplotlib import cm

#@profile
def Method_Iterative(A, b, Precision):

	Iteration_Matrix, Constant_Matrix, Initial_Vector = Generate_Matrix_for_Iterate(A, b)
	Error = abs(Precision) + 1
	X = Initial_Vector
	conut = 0
	while Error > abs(Precision):
		X_next = Iteration_Matrix*X + Constant_Matrix
		Error = abs(np.mean(X_next - X))
		X = X_next
		conut = conut + 1
		#print(Error)
	return X, Error, conut

def Generate_Matrix_for_Iterate(A, b):
	'''
	A = M - N 
	N = M - A
	x = M^{-1}*N*x + M^{-1}*b
	x = Iteration_Matrix*x + Constant_Matrix
	'''
	Dimention = A.shape[0]
	M = A.item(0,0)*np.matrix(np.eye(Dimention,Dimention))
	N = M - A
	M_inv = 1/A.item(0,0)#*np.matrix(np.eye(Dimention,Dimention))
	Iteration_Matrix = M_inv*N#np.dot(M_inv,N)#
	Constant_Matrix = M_inv*b
	Initial_Vector = 0*np.matrix(np.ones([Dimention,1]))
	#print(Iteration_Matrix)
	return Iteration_Matrix, Constant_Matrix, Initial_Vector

def calGCD(op1, op2):
	if (op2==0): 
		return op1
	else: 
		return calGCD(op2, op1%op2)

def calLCM(op1, op2):
	gcd = calGCD(op1, op2)
	lcm = op1/gcd*op2
	return lcm

def Generate_Matrix_A_b(Nx, Ny):
	#Diagonal Matrix Kernal
	#D = np.matrix(Diagonal_Matrix_String)
	#Kernal_Dim = D.shape[0]
	#The Matrix near by Diagonal Matrix Kernal
	#H = np.matrix(np.eye(Kernal_Dim))
	#The Zero Matrix
	#Z = np.matrix(np.zeros([Kernal_Dim,Kernal_Dim]))
	#print(D, H, Z)
	#Kernel_String = "H,D,H,Z,Z"
	D, H, Z, LCM = Generate_Diagonal_Matrix(Nx, Ny)
	Kernel_String = "H,D,H,Z"
	Some_Z = ",Z"*(Nx-3)
	Kernel_String = Kernel_String + Some_Z
	Matrix_String = Kernel_String[2:]
	#First -1 because Kernel_String is one more than real, Second -1 becasue there is a value out of loop
	for i in range(int((len(Kernel_String)-1)/2)-1):
		Matrix_String = Matrix_String+";"+"Z,"*i+Kernel_String[0:len(Kernel_String)-(i+1)*2]
	A = np.bmat(Matrix_String)

			# Vector_b_String = str(int(np.random.random()*10))
			# #First -1 because Kernel_String is one more than real, Second -1 becasue there is a value out of loop
			# for i in range(int((len(Kernel_String)-1)/2*D.shape[0])-1):
			# 	Vector_b_String = Vector_b_String+";"+str(int(np.random.random()*10))
			# b = np.matrix(Vector_b_String)
			# with open("matrix_file", 'w') as mf:
			# 	A.tofile(mf)
			# 	b.tofile(mf)
			# np.savetxt("matrix_file", A)
	b = np.zeros(Nx*Ny)
	Border_Up, Border_Down, Border_Left, Border_Right = Border_Generator(Nx, Ny)
	for i in range(Nx):
		for j in range(Ny):
			if 0 == i:
				b[Ny*i+j] = b[Ny*i+j] + (-LCM/Nx**2)*Border_Left[j+1]
			if (Nx-1) == i:
				b[Ny*i+j] = b[Ny*i+j] + (-LCM/Nx**2)*Border_Right[j+1]
			if 0 == j:
				b[Ny*i+j] = b[Ny*i+j] + (-LCM/Ny**2)*Border_Up[i]
			if (Ny-1) == j:
				b[Ny*i+j] = b[Ny*i+j] + (-LCM/Ny**2)*Border_Down[i]
	b = np.mat(b).T
	return A, b

# A = np.matrix("-4,1,0,1,0,0,0,0,0;1,-4,1,0,1,0,0,0,0;0,1,-4,0,0,1,0,0,0;1,0,0,-4,1,0,1,0,0;0,1,0,1,-4,1,0,1,0;0,0,1,0,1,-4,0,0,1;0,0,0,1,0,0,-4,1,0;0,0,0,0,1,0,1,-4,1;0,0,0,0,0,1,0,1,-4")
# b=np.matrix("5;6;7;8;9;7;8;9;2")
def Generate_Diagonal_Matrix(Nx, Ny):
	"""
	Nx和Ny分别表示了在X方向和Y方向上划分的网格数
	"""
	#Block_Line_String = "1,-4,1,0"
	#首先确定X方向和Y方向上划分的网格数平方的最小公倍数，以确定系数
	LCM = calLCM(Nx**2, Ny**2)
	Block_Line_String = str(int(LCM/Ny**2))+","+str(int((-2)*(LCM/Nx**2+LCM/Ny**2)))+","+str(int(LCM/Ny**2))+",0"
	#-3是因为在上一行已经产生了4个元素，而第1个元素是为了生成方便加上的，因此在上一行产生了3个元素
	Block_Line_String = Block_Line_String + ",0"*(Ny-3)
	index = Block_Line_String.find(",")+1
	Diagonal_Matrix_String = Block_Line_String[index:]
	position = Block_Line_String.rfind(",")
	for i in range(Block_Line_String.count(",")-1):
		Diagonal_Matrix_String = Diagonal_Matrix_String+";"+"0,"*i+Block_Line_String[0:position]
		position = Block_Line_String.rfind(",",0,position)
	#print(Diagonal_Matrix_String)

	D = np.matrix(Diagonal_Matrix_String)
	Kernal_Dim = D.shape[0]
	H = np.matrix(np.eye(Kernal_Dim))*int(LCM/Nx**2)
	Z = np.matrix(np.zeros([Kernal_Dim,Kernal_Dim]))
	return D, H, Z, LCM


def Border_Generator(Nx, Ny):
	Border_Up = np.linspace(0,0,Nx)
	Border_Down = np.linspace(5,5,Nx)
	Border_Left = np.linspace(0,5,Ny+2)
	Border_Right = np.linspace(0,5,Ny+2)
	print(Border_Up, Border_Down, Border_Left, Border_Right)
	return Border_Up, Border_Down, Border_Left, Border_Right

def Plot_Result(Z, Nx, Ny, Xmax, Xmin, Ymax, Ymin):
	xs = np.linspace(Xmin, Xmax, Nx)
	ys = np.linspace(Ymin, Ymax, Ny)
	X, Y = np.meshgrid(xs, ys)
	Z = np.array(np.reshape(Z,[Nx,Ny]).T)
	fig = plt.figure()
	ax = fig.gca(projection='3d')

	cset = ax.contour(X, Y, Z, zdir='z', offset=0, cmap=cm.coolwarm)
	cset = ax.contour(X, Y, Z, zdir='x', offset=6, cmap=cm.coolwarm)
	cset = ax.contour(X, Y, Z, zdir='y', offset=6, cmap=cm.coolwarm)
	ax.plot_surface(X, Y, Z, rstride=8, cstride=8, alpha=0.3)
	ax.set_xlabel('X')
	ax.set_xlim(0, 6)
	ax.set_ylabel('Y')
	ax.set_ylim(0, 6)
	ax.set_zlabel('Z')
	ax.set_zlim(0, 6)
	plt.show()

if __name__ == "__main__":
	#np.set_printoptions(threshold='nan')
	Nx = 30
	Ny = 30
	A, b = Generate_Matrix_A_b(Nx, Ny)
	print(b)
	# start_time = clock()
	# print("Direct Method:\n", A.I*b, A.shape)
	# end_time = clock()
	# print("\nUse time:", end_time - start_time)

	print("Begin Calcule!\n")
	start_time = clock()
	Z, Error, count = Method_Iterative(A, b, 1e-5)
	print("\nIterative Method:\n", Z, Error, count)
	end_time = clock()
	print("\nUse time:", end_time - start_time)

	Plot_Result(Z, Nx, Ny, Xmax=5, Xmin=0, Ymax=5, Ymin=0)

	# #Diagonal Matrix Kernal
	# D = np.matrix("-4,1,0;1,-4,0;0,1,-4")
	# #The Matrix near by Diagonal Matrix Kernal
	# H = np.matrix("1,0,0;0,1,0;0,0,1")
	# #The Zero Matrix
	# Z = np.matrix("0,0,0;0,0,0;0,0,0")

	# Kernel_String="H,D,H,Z,Z,Z,Z"
	# Matrix_String=Kernel_String[2:]
	# for i in range(int((len(Kernel_String)-1)/2)-1):
	# 	Matrix_String = Matrix_String+";"+"Z,"*i+Kernel_String[0:len(Kernel_String)-(i+1)*2]
	# A = np.bmat(Matrix_String)
